# -*- coding: utf-8 -*-
"""
evacuation_dynamic_exit.py

动态出口选择的三维疏散模拟，结合体力消耗与拥堵惩罚
策略：
1. 为每个出口分别计算静态场(static field)
2. 每个Agent动态评估到各出口的代价：
   cost = 距离 + beta * 当前目标出口等待队列长度
3. 选择最优出口，并沿该出口的静态场梯度下降
4. 邻域移动时，考虑体力消耗调整步长
5. 可视化时，不同目标出口的Agent用不同颜色标记

注释均为中文
"""
import numpy as np
import heapq
import math
import random
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import animation
from scipy.ndimage import binary_dilation
from collections import deque

# -----------------------------
# 1. 参数设置
# -----------------------------
CELL_SIZE          = 1.0     # 元胞边长 (m)
FLOORS             = 10      # 楼层数
HEIGHT, WIDTH      = 30, 30  # 每层网格尺寸
AGENTS_PER_FLOOR   = 40      # 每层初始人数
SPEED_MEAN         = 1.0     # m/s，初始速度均值
SPEED_STD          = 0.1     # m/s，速度标准差
STAIR_SPEED_FACTOR = 0.7     # 楼梯移动速度系数
FLOOR_HEIGHT       = 3.0     # 每层楼高 (m)
MAX_STEPS          = 1000    # 最大模拟步数
STAIR_COUNT        = 3       # 每层楼梯数量
OBSTACLE_DENSITY   = 0.25    # 障碍物密度
ANIMATION_SPEED    = 500     # 动画间隔 (ms)
REASSIGN_INTERVAL  = 20      # 动态重选出口的步数间隔
BETA               = 5.0     # 拥堵惩罚系数

# -----------------------------
# 2. 迷宫与环境生成（保持原逻辑）
# -----------------------------
def generate_maze(height, width):
    cell_size = 2
    maze = np.ones((height, width), dtype=int)
    visited = np.zeros((height//cell_size, width//cell_size), dtype=bool)
    def carve(y, x):
        dirs = [(-1,0),(1,0),(0,-1),(0,1)]
        random.shuffle(dirs)
        for dy, dx in dirs:
            ny, nx = y+dy, x+dx
            if 0 <= ny < height//cell_size and 0 <= nx < width//cell_size and not visited[ny,nx]:
                visited[ny,nx] = True
                cy, cx = y*cell_size, x*cell_size
                ny0, nx0 = ny*cell_size, nx*cell_size
                if dy != 0:
                    for i in range(cell_size//2):
                        maze[cy + dy*(cell_size//2 - i)][cx] = 0
                        maze[ny0 - dy*(i+1)][nx0] = 0
                else:
                    for i in range(cell_size//2):
                        maze[cy][cx + dx*(cell_size//2 - i)] = 0
                        maze[ny0][nx0 - dx*(i+1)] = 0
                carve(ny, nx)
    visited[0,0] = True
    carve(0,0)
    # 随机扩宽通道
    for _ in range(random.randint(1,3)):
        struct = random.choice([
            np.array([[0,1,0],[0,1,0],[0,1,0]]),
            np.array([[0,0,0],[1,1,1],[0,0,0]]),
            np.ones((3,3))
        ])
        maze = binary_dilation(maze==0, structure=struct).astype(int)
    maze[1,1] = 0
    maze[-2,-2] = 0
    return maze

def build_environment():
    env = np.ones((FLOORS, HEIGHT, WIDTH), dtype=int)
    stair_positions = []
    for z in range(FLOORS):
        maze = generate_maze(HEIGHT, WIDTH)
        sp = []
        for _ in range(STAIR_COUNT):
            while True:
                y, x = random.randint(2,HEIGHT-3), random.randint(2,WIDTH-3)
                if maze[y,x] == 0:
                    sp.append((y,x)); break
        for (y,x) in sp:
            maze[y-1:y+2, x-1:x+2] = 0
            maze[y,x] = 2
        env[z] = maze
        stair_positions.append(sp)
    # 出口集合（所有底层楼梯即为出口）
    exits = [(0, y, x) for y,x in stair_positions[0]]
    return env, stair_positions, exits

# -----------------------------
# 3. 静态场计算（每出口独立）
# -----------------------------
def compute_static_field(env, exits):
    f,h,w = env.shape
    static = np.full((f,h,w), np.inf)
    vis = np.zeros((f,h,w), dtype=bool)
    pq = []
    for e in exits:
        static[e] = 0.0
        heapq.heappush(pq, (0.0, e))
    neigh = [(dz,dy,dx) for dz in (-1,0,1) for dy in (-1,0,1) for dx in (-1,0,1) if not (dz==dy==dx==0)]
    while pq:
        d,(z,y,x) = heapq.heappop(pq)
        if vis[z,y,x]: continue
        vis[z,y,x] = True
        for dz,dy,dx in neigh:
            nz,ny,nx = z+dz, y+dy, x+dx
            if 0<=nz<f and 0<=ny<h and 0<=nx<w and env[nz,ny,nx]!=1:
                step = CELL_SIZE * (math.sqrt(1+(FLOOR_HEIGHT/CELL_SIZE)**2) if dz!=0 else 1)
                nd = d + step
                if nd < static[nz,ny,nx]:
                    static[nz,ny,nx] = nd
                    heapq.heappush(pq, (nd, (nz,ny,nx)))
    return static

def compute_all_static(env, exits):
    # 为每个出口计算静态场
    fields = []
    for e in exits:
        fields.append(compute_static_field(env, [e]))
    return np.stack(fields)

# -----------------------------
# 4. Agent 类，带目标出口和体力
# -----------------------------
class Agent:
    def __init__(self, z,y,x, static_fields, queue_counts):
        self.pos = (z,y,x)
        self.V0 = random.gauss(SPEED_MEAN, SPEED_STD)
        self.d_sum = 0.0
        self.t_sum = 0.0
        self.counter = 0
        self.static_fields = static_fields
        self.queue_counts = queue_counts
        # 初始化分配最近出口
        costs = [static_fields[i][self.pos] + BETA*queue_counts[i] for i in range(len(queue_counts))]
        self.target = int(np.argmin(costs))
        queue_counts[self.target] += 1

    def speed(self, dz):
        base = self.V0 * (1 - 0.01*(self.d_sum**0.71)/(1 - 0.01*(self.t_sum**0.51)*math.exp(0.01*self.t_sum)))
        return base * (STAIR_SPEED_FACTOR if dz!=0 else 1.0)

    def reassign(self):
        # 动态评估各出口代价，重新分配
        old = self.target
        self.queue_counts[old] = max(0, self.queue_counts[old]-1)
        costs = [self.static_fields[i][self.pos] + BETA*self.queue_counts[i] for i in range(len(self.queue_counts))]
        self.target = int(np.argmin(costs))
        self.queue_counts[self.target] += 1

# -----------------------------
# 5. 初始化 Agents
# -----------------------------
def init_agents(env, static_fields):
    Nexit = static_fields.shape[0]
    queue_counts = [0]*Nexit
    agents = []
    for z in range(FLOORS):
        empties = list(zip(*np.where(env[z]==0)))
        random.shuffle(empties)
        for y,x in empties[:AGENTS_PER_FLOOR]:
            agents.append(Agent(z,y,x, static_fields, queue_counts))
    return agents, queue_counts

# -----------------------------
# 6. 模拟循环
# -----------------------------
def simulate(env, static_fields, exits):
    agents, queue_counts = init_agents(env, static_fields)
    frames, stats = [], []
    neigh = [(dz,dy,dx) for dz in (-1,0,1) for dy in (-1,0,1) for dx in (-1,0,1) if not (dz==dy==dx==0)]

    for step in range(MAX_STEPS):
        occ = set(a.pos for a in agents)
        positions=[]
        # 动态重选出口
        if step>0 and step % REASSIGN_INTERVAL == 0:
            for a in agents:
                a.reassign()

        for a in agents[:]:
            a.t_sum += 1; a.counter += 1
            z,y,x = a.pos
            tau = CELL_SIZE / a.speed(0)
            if a.counter>=tau:
                cands=[]
                sf = static_fields[a.target]
                for dz,dy,dx in neigh:
                    nz,ny,nx = z+dz,y+dy,x+dx
                    if 0<=nz<FLOORS and 0<=ny<HEIGHT and 0<=nx<WIDTH and env[nz,ny,nx]!=1 and (nz,ny,nx) not in occ:
                        if sf[nz,ny,nx] <= sf[z,y,x]:
                            cands.append((sf[nz,ny,nx], (dz,dy,dx)))
                if cands:
                    cands.sort(key=lambda x: x[0])
                    best = [mv for d,mv in cands if abs(d - cands[0][0])<1e-6]
                    dz,dy,dx = random.choice(best)
                    new = (z+dz,y+dy,x+dx)
                    occ.remove(a.pos); occ.add(new)
                    a.pos=new
                    dist = CELL_SIZE*(math.sqrt(1+(FLOOR_HEIGHT/CELL_SIZE)**2) if dz!=0 else 1)
                    a.d_sum += dist; a.counter=0
                    # 若到达最终出口，则移除Agent
                    if new in exits:
                        occ.remove(new)
                        queue_counts[a.target] = max(0, queue_counts[a.target]-1)
                        agents.remove(a)
            positions.append((a.pos, a.target))
        stats.append((len(agents), [sum(1 for _,t in positions if t==i) for i in range(len(exits))]))
        frames.append(positions)
        if not agents: break
    return frames, stats

# -----------------------------
# 7. 可视化
# -----------------------------
def visualize(env, frames, stats, stair_positions, exits):
    cmap = ListedColormap(['white','dimgray','saddlebrown'])
    exit_colors = ['blue','green','red','orange','purple']
    fig = plt.figure(figsize=(20,10))
    gs = fig.add_gridspec(2,6)
    ax3d = fig.add_subplot(gs[:,0], projection='3d')
    ax3d.set_box_aspect([WIDTH,HEIGHT,FLOORS])
    ax3d.set_xlabel('X'); ax3d.set_ylabel('Y'); ax3d.set_zlabel('Floor')
    stairs = (env==2)
    ax3d.voxels(stairs, facecolors='saddlebrown', edgecolor='k', alpha=0.7)
    scat3d = ax3d.scatter([],[],[], s=20)
    txt = ax3d.text2D(0.02,0.95,'', transform=ax3d.transAxes)

    axs, scats = [], []
    for f in range(FLOORS):
        r,c = (0 if f<5 else 1), (f%5)+1
        ax = fig.add_subplot(gs[r,c])
        ax.imshow(env[f], cmap=cmap, origin='lower')
        ax.set_title(f'Floor {f+1}'); ax.axis('off')
        for sy,sx in stair_positions[f]: ax.scatter(sx,sy,c='red',s=50,marker='s')
        axs.append(ax)
        scats.append(ax.scatter([],[],s=10))

    def update(i):
        pts = frames[i]
        xs,ys,zs,cs=[],[],[],[]
        for (z,y,x),t in pts:
            xs.append(x); ys.append(y); zs.append(z)
            cs.append(exit_colors[t])
        scat3d._offsets3d = (xs,ys,zs)
        scat3d.set_color(cs)
        txt.set_text(f'剩余: {len(pts)}')
        for f in range(FLOORS):
            floor_pts = [(y,x) for (z,y,x),t in pts if z==f]
            if floor_pts:
                ys_f, xs_f = zip(*floor_pts)
                scats[f].set_offsets(np.column_stack([xs_f, ys_f]))
                scats[f].set_color([cs[j] for j,(pos,_) in enumerate(pts) if pos[0]==f])
            else:
                scats[f].set_offsets(np.empty((0,2)))
        return [scat3d] + scats + [txt]

    anim = animation.FuncAnimation(fig, update, frames=len(frames), interval=ANIMATION_SPEED, blit=False)
    plt.show()

# -----------------------------
# 8. 主程序
# -----------------------------
def main():
    env, stair_positions, exits = build_environment()
    static_fields = compute_all_static(env, exits)
    frames, stats = simulate(env, static_fields, exits)
    visualize(env, frames, stats, stair_positions, exits)

if __name__ == '__main__':
    main()
